{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from llama_index.core import SimpleDirectoryReader\n",
    "import ollama\n",
    "from groq import Groq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"GROQ_API_KEY\"] = \"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\"\n",
    "# os.environ[\"GROQ_API_KEY\"] = \"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\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "\n",
    "def load_documents(path: str):\n",
    "    # reader = SimpleDirectoryReader(input_dir=path, recursive=True)\n",
    "    reader = SimpleDirectoryReader(input_dir=path)\n",
    "    documents = reader.load_data()\n",
    "    doc_dicts = [{\"content\": d.text, **d.metadata} for d in documents]\n",
    "    return doc_dicts\n",
    "\n",
    "\n",
    "def process_metadata(doc_dicts):\n",
    "    file_seen = set()\n",
    "    metadata_list = []\n",
    "    for doc in doc_dicts:\n",
    "        if doc[\"file_path\"] not in file_seen:\n",
    "            file_seen.add(doc[\"file_path\"])\n",
    "            metadata_list.append(doc)\n",
    "    return metadata_list\n",
    "\n",
    "\n",
    "def query_summaries(doc_dicts):\n",
    "    client = Groq(\n",
    "        api_key=os.environ.get(\"GROQ_API_KEY\"),\n",
    "    )\n",
    "    \n",
    "    PROMPT = f\"\"\"\n",
    "    The following is a list of file contents, along with their metadata. For each file, provide a summary of the contents.\n",
    "    The purpose of the summary is to organize files based on their content. To this end provide a concise but informative summary.\n",
    "    Try to make the summary as specific to the file as possible.\n",
    "\n",
    "    {doc_dicts}\n",
    "\n",
    "    Return a JSON list with the following schema:\n",
    "\n",
    "    ```json\n",
    "    {{\n",
    "    \"files\": [\n",
    "        {{\n",
    "        \"file_path\": \"path to the file including name\",\n",
    "        \"summary\": \"summary of the content\"\n",
    "        }}\n",
    "    ]\n",
    "    }}\n",
    "    ```\n",
    "    \"\"\".strip()\n",
    "    \n",
    "\n",
    "    chat_completion = client.chat.completions.create(\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": \"Always return JSON. Do not include any other text or formatting characters.\",\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": PROMPT,\n",
    "            },\n",
    "        ],\n",
    "        model=\"llama3-70b-8192\",\n",
    "        response_format={\"type\": \"json_object\"},\n",
    "    )\n",
    "\n",
    "    summaries = json.loads(chat_completion.choices[0].message.content)[\"files\"]\n",
    "\n",
    "    return summaries\n",
    "\n",
    "\n",
    "def merge_summary_documents(summaries, metadata_list):\n",
    "    list_summaries = defaultdict(list)\n",
    "\n",
    "    for item in summaries:\n",
    "        list_summaries[item[\"file_path\"]].append(item[\"summary\"])\n",
    "\n",
    "    file_summaries = {path: '. '.join(summaries) for path, summaries in list_summaries.items()}\n",
    "\n",
    "    file_list = [{\"summary\": file_summaries[file[\"file_path\"]], **file} for file in metadata_list]\n",
    "    \n",
    "    return file_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_doc_summaries(path: str):# -> List[Dict[str, str]]:\n",
    "    doc_dicts = load_documents(path)\n",
    "    metadata = process_metadata(doc_dicts)\n",
    "    \n",
    "    summaries = query_summaries(doc_dicts)\n",
    "    \n",
    "    file_summaries = merge_summary_documents(summaries, metadata)\n",
    "    \n",
    "    return file_summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_summaries = get_doc_summaries(\"../\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'summary': 'MIT License file containing copyright and permission notice.',\n",
       "  'content': 'MIT License\\n\\nCopyright (c) 2024 Ajay Arasanipalai\\n\\nPermission is hereby granted, free of charge, to any person obtaining a copy\\nof this software and associated documentation files (the \"Software\"), to deal\\nin the Software without restriction, including without limitation the rights\\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\\ncopies of the Software, and to permit persons to whom the Software is\\nfurnished to do so, subject to the following conditions:\\n\\nThe above copyright notice and this permission notice shall be included in all\\ncopies or substantial portions of the Software.\\n\\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\\nSOFTWARE.\\n',\n",
       "  'file_path': '/Users/ashwin/side_projects/llama-fs/notebooks/../LICENSE',\n",
       "  'file_name': 'LICENSE',\n",
       "  'file_size': 1074,\n",
       "  'creation_date': '2024-05-11',\n",
       "  'last_modified_date': '2024-05-11'},\n",
       " {'summary': 'Jupyter notebook with various code snippets related to Groq and llama-index.. Jupyter notebook with code related to environment variable and API key.. Jupyter notebook with code related to document processing and metadata.. Jupyter notebook with code related to document processing and JSON output.. Jupyter notebook with code related to Groq and llama-index integration.. Jupyter notebook with code related to document processing and Groq API.',\n",
       "  'content': '\\n\\n\\nimport openai\\nimport groq\\nimport os\\nimport json\\nfrom groq import Groq\\n\\n\\n# ',\n",
       "  'file_path': '/Users/ashwin/side_projects/llama-fs/notebooks/../XXXXXX.ipynb',\n",
       "  'file_name': 'XXXXXX.ipynb',\n",
       "  'file_size': 133781,\n",
       "  'creation_date': '2024-05-11',\n",
       "  'last_modified_date': '2024-05-11'},\n",
       " {'summary': 'Python script with various functions for document processing and summary generation.',\n",
       "  'content': 'import json\\nimport os\\nfrom collections import defaultdict\\nfrom llama_index.core import SimpleDirectoryReader\\nimport ollama\\nfrom groq import Groq\\n\\n\\ndef get_doc_summaries(path: str) -> List[Dict[str, str, str, str, str, str]]:\\n    doc_dicts = load_documents(path)\\n    metadata = process_metadata(doc_dicts)\\n    \\n    summaries = query_summaries(doc_dicts)\\n    \\n    file_summaries = merge_summary_documents(summaries, metadata)\\n    \\n    return file_summaries\\n\\n    # [\\n    #     {\\n    #         file_path: \\n    #         file_name:\\n    #         file_size: \\n    #         content: \\n    #         summary: \\n    #         creation_date:\\n    #         last_modified_date:\\n    #     }\\n    # ]\\n\\n\\ndef load_documents(path: str):\\n    # reader = SimpleDirectoryReader(input_dir=path, recursive=True)\\n    reader = SimpleDirectoryReader(input_dir=path)\\n    documents = reader.load_data()\\n    doc_dicts = [{\"content\": d.text, **d.metadata} for d in documents]\\n    return doc_dicts\\n\\n\\ndef process_metadata(doc_dicts):\\n    file_seen = set()\\n    metadata_list = []\\n    for doc in doc_dicts:\\n        if doc[\"file_path\"] not in file_seen:\\n            file_seen.add(doc[\"file_path\"])\\n            metadata_list.append(doc)\\n    return metadata_list\\n\\n\\ndef query_summaries(doc_dicts):\\n    client = Groq(\\n        api_key=os.environ.get(\"GROQ_API_KEY\"),\\n    )\\n    \\n    PROMPT = f\"\"\"\\n    The following is a list of file contents, along with their metadata. For each file, provide a summary of the contents.\\n\\n    {doc_dicts}\\n\\n    Return a JSON list with the following schema:\\n\\n    ```json\\n    {{\\n    \"files\": [\\n        {{\\n        \"file_path\": \"path to the file including name\",\\n        \"summary\": \"summary of the content\"\\n        }}\\n    ]\\n    }}\\n    ```\\n    \"\"\".strip()\\n    \\n\\n    chat_completion = client.chat.completions.create(\\n        messages=[\\n            {\\n                \"role\": \"system\",\\n                \"content\": \"Always return JSON. Do not include any other text or formatting characters.\",\\n            },\\n            {\\n                \"role\": \"user\",\\n                \"content\": PROMPT,\\n            },\\n        ],\\n        model=\"llama3-70b-8192\",\\n        response_format={\"type\": \"json_object\"},\\n    )\\n\\n    summaries = json.loads(chat_completion.choices[0].message.content)[\"files\"]\\n\\n    return summaries\\n\\n\\ndef merge_summary_documents(summaries, metadata_list):\\n    list_summaries = defaultdict(list)\\n\\n    for item in summaries:\\n        list_summaries[item[\"file_path\"]].append(item[\"summary\"])\\n\\n    file_summaries = {path: \\'. \\'.join(summaries) for path, summaries in list_summaries.items()}\\n\\n    file_list = [{\"summary\": file_summaries[file[\"file_path\"]], **file} for file in metadata_list]\\n    \\n    return file_list',\n",
       "  'file_path': '/Users/ashwin/side_projects/llama-fs/notebooks/../loader.py',\n",
       "  'file_name': 'loader.py',\n",
       "  'file_type': 'text/x-python',\n",
       "  'file_size': 2696,\n",
       "  'creation_date': '2024-05-11',\n",
       "  'last_modified_date': '2024-05-11'},\n",
       " {'summary': 'List of required Python packages for the project.',\n",
       "  'content': 'ollama\\nchromadb\\nllama-index\\nlitellm\\ngroq\\ndocx2txt',\n",
       "  'file_path': '/Users/ashwin/side_projects/llama-fs/notebooks/../requirements.txt',\n",
       "  'file_name': 'requirements.txt',\n",
       "  'file_type': 'text/plain',\n",
       "  'file_size': 49,\n",
       "  'creation_date': '2024-05-11',\n",
       "  'last_modified_date': '2024-05-11'}]"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file_summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama3_hackathon",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
